package com.ocean.web.util;

import ucar.ma2.*;
import ucar.nc2.*;

import java.io.File;
import java.io.IOException;
import java.util.*;

/**
 * 在分析产品
 */
public class CpjypgZFXCP {
    // 裁剪范围（可根据需求改为参数传入）
    private static final double MIN_LAT = 14.0;
    private static final double MAX_LAT = 42.0;
    private static final double MIN_LON = 105.0;
    private static final double MAX_LON = 132.0;

    public static void main(String[] args) {
        try {
            // 1. 定义输入文件配置（作为参数，可外部传入）
            Map<String, String> swhFiles = new HashMap<String, String>() {{
                put("E:\\项目资料备份\\新云平台文件\\2025工作\\再分析产品\\wave\\ERA_SWH_20250908.nc", "swh");
            }};
            Map<String, String> windFiles = new HashMap<String, String>() {{
                put("E:\\项目资料备份\\新云平台文件\\2025工作\\再分析产品\\wind\\ERA_U10_20250908.nc", "u10");
                put("E:\\项目资料备份\\新云平台文件\\2025工作\\再分析产品\\wind\\ERA_V10_20250908.nc", "v10");
            }};
            Map<String, String> sstFiles = new HashMap<String, String>() {{
                put("E:\\项目资料备份\\新云平台文件\\2025工作\\再分析产品\\sst\\ERA_SST_20250908.nc", "sst");
            }};

            // 2. 定义输出文件路径（也可作为参数传入）
            String swhOutput = "E:\\swh_cropped.nc";
            String windOutput = "E:\\wind_calculated.nc";
            String sstOutput = "E:\\sst_cropped.nc";

            // 3. 调用处理方法，传入输入配置和输出路径
            processSwh(swhFiles, swhOutput);
            processWind(windFiles, windOutput);
            processSst(sstFiles, sstOutput);

            System.out.println("\n所有存在的文件处理完成！");
        } catch (Exception e) {
            System.err.println("处理失败：" + e.getMessage());
            e.printStackTrace();
        }
    }

    // ==================== SWH处理 ====================
    /**
     * 处理SWH文件
     * @param fileMap  输入文件配置（路径->变量名）
     * @param outputPath 输出文件路径
     */
    public static void processSwh(Map<String, String> fileMap, String outputPath) throws Exception {
        processFileGroup("SWH", fileMap, () -> {
            try {
                innerProcessSwh(fileMap, outputPath);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }

    private static void innerProcessSwh(Map<String, String> fileMap, String outputPath) throws Exception {
        Map<String, Array> croppedDataMap = new HashMap<>();
        double[] lats = null;
        double[] lons = null;
        int timeSize = -1;
        int latSize = -1;
        int lonSize = -1;
        NetcdfFile refNc = null;
        String validVarName = null;

        for (Map.Entry<String, String> entry : fileMap.entrySet()) {
            String filePath = entry.getKey();
            String varName = entry.getValue();

            if (!new File(filePath).exists()) {
                System.out.println("文件不存在，跳过：" + filePath);
                continue;
            }

            System.out.println("处理SWH文件：" + filePath + "，变量：" + varName);
            validVarName = varName;

            NetcdfFile ncFile = NetcdfFile.open(filePath);
            if (refNc == null) refNc = ncFile;

            Variable var = ncFile.findVariable(varName);
            Variable latVar = ncFile.findVariable("latitude");
            Variable lonVar = ncFile.findVariable("longitude");
            validateVariables(var, latVar, lonVar, varName);

            if (lats == null || lons == null) {
                lats = readLatLonArray(latVar);
                lons = readLatLonArray(lonVar);
                int[] latIndices = getRangeIndices(lats, MIN_LAT, MAX_LAT);
                int[] lonIndices = getRangeIndices(lons, MIN_LON, MAX_LON);
                if (latIndices == null || lonIndices == null) {
                    throw new RuntimeException("裁剪范围超出数据范围：" + filePath);
                }
                latSize = latIndices[1];
                lonSize = lonIndices[1];
                System.out.printf("SWH裁剪后维度：纬度%d个点，经度%d个点%n", latSize, lonSize);
            }

            int[] shape = var.getShape();
            if (shape.length != 3) throw new RuntimeException(varName + "必须是3维变量（time, lat, lon）");
            if (timeSize == -1) timeSize = shape[0];

            int[] origin = new int[]{0, getRangeIndices(lats, MIN_LAT, MAX_LAT)[0], getRangeIndices(lons, MIN_LON, MAX_LON)[0]};
            int[] cropShape = new int[]{timeSize, latSize, lonSize};
            Array croppedArray = cropVariable(var, origin, cropShape);
            croppedDataMap.put(varName, croppedArray);

            ncFile.close();
            System.out.println("SWH裁剪完成");
        }

        if (refNc != null && validVarName != null) {
            createSingleVarNcFile(refNc, croppedDataMap, lats, lons, timeSize, latSize, lonSize, outputPath, validVarName, true);
        } else {
            System.out.println("SWH无有效文件，不生成输出");
        }
    }

    // ==================== 风场处理 ====================
    /**
     * 处理风场文件
     * @param fileMap  输入文件配置（路径->变量名，需包含u10和v10）
     * @param outputPath 输出文件路径
     */
    public static void processWind(Map<String, String> fileMap, String outputPath) throws Exception {
        processFileGroup("风场", fileMap, () -> {
            try {
                innerProcessWind(fileMap, outputPath);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }

    private static void innerProcessWind(Map<String, String> fileMap, String outputPath) throws Exception {
        // 校验u10和v10是否都存在
        String u10Path = null, v10Path = null;
        for (Map.Entry<String, String> entry : fileMap.entrySet()) {
            if (entry.getValue().equals("u10")) u10Path = entry.getKey();
            if (entry.getValue().equals("v10")) v10Path = entry.getKey();
        }
        if (u10Path == null || v10Path == null || !new File(u10Path).exists() || !new File(v10Path).exists()) {
            throw new RuntimeException("风场文件u10或v10不存在，无法处理");
        }

        Map<String, Array> croppedDataMap = new HashMap<>();
        Map<String, Variable> windVarMetadata = new HashMap<>();
        double[] lats = null;
        double[] lons = null;
        Array croppedLatsArray = null;
        Array croppedLonsArray = null;
        int timeSize = -1;
        int latSize = -1;
        int lonSize = -1;
        NetcdfFile firstWindNc = null;

        for (Map.Entry<String, String> entry : fileMap.entrySet()) {
            String filePath = entry.getKey();
            String varName = entry.getValue();

            if (!new File(filePath).exists()) {
                System.out.println("文件不存在，跳过：" + filePath);
                continue;
            }

            System.out.println("处理风场文件：" + filePath + "，变量：" + varName);
            NetcdfFile ncFile = NetcdfFile.open(filePath);
            if (firstWindNc == null) firstWindNc = ncFile;

            Variable var = ncFile.findVariable(varName);
            Variable latVar = ncFile.findVariable("latitude");
            Variable lonVar = ncFile.findVariable("longitude");
            validateVariables(var, latVar, lonVar, varName);
            windVarMetadata.put(varName, var);

            if (lats == null || lons == null) {
                lats = readLatLonArray(latVar);
                lons = readLatLonArray(lonVar);
                int[] latIndices = getRangeIndices(lats, MIN_LAT, MAX_LAT);
                int[] lonIndices = getRangeIndices(lons, MIN_LON, MAX_LON);
                if (latIndices == null || lonIndices == null) {
                    throw new RuntimeException("裁剪范围超出数据范围：" + filePath);
                }
                latSize = latIndices[1];
                lonSize = lonIndices[1];
                System.out.printf("风场裁剪后维度：纬度%d个点，经度%d个点%n", latSize, lonSize);
                croppedLatsArray = cropVariable(latVar, new int[]{latIndices[0]}, new int[]{latSize});
                croppedLonsArray = cropVariable(lonVar, new int[]{lonIndices[0]}, new int[]{lonSize});
            }

            int[] shape = var.getShape();
            if (shape.length != 3) throw new RuntimeException(varName + "必须是3维变量（time, lat, lon）");
            if (timeSize == -1) timeSize = shape[0];
            else if (shape[0] != timeSize) throw new RuntimeException("u10/v10时间维度不一致");

            int[] origin = new int[]{0, getRangeIndices(lats, MIN_LAT, MAX_LAT)[0], getRangeIndices(lons, MIN_LON, MAX_LON)[0]};
            int[] cropShape = new int[]{timeSize, latSize, lonSize};
            Array croppedArray = cropVariable(var, origin, cropShape);
            croppedDataMap.put(varName, croppedArray);

            ncFile.close();
            System.out.println(varName + "裁剪完成");
        }

        // 计算风速风向
        Array u10Array = croppedDataMap.get("u10");
        Array v10Array = croppedDataMap.get("v10");
        if (u10Array == null || v10Array == null) throw new RuntimeException("u10或v10数据缺失");

        float[] u10Data = (float[]) u10Array.copyTo1DJavaArray();
        float[] v10Data = (float[]) v10Array.copyTo1DJavaArray();
        int dataSize = u10Data.length;

        float[] wspd10Data = new float[dataSize];
        float[] wdir10Data = new float[dataSize];
        for (int i = 0; i < dataSize; i++) {
            float u = u10Data[i];
            float v = v10Data[i];
            wspd10Data[i] = (float) Math.sqrt(u * u + v * v);
            double rad = Math.atan2(u, v);
            wdir10Data[i] = (float) (rad * 180 / Math.PI + 180);
        }

        int[] dataShape = new int[]{timeSize, latSize, lonSize};
        croppedDataMap.put("wspd10", Array.factory(DataType.FLOAT, dataShape, wspd10Data));
        croppedDataMap.put("wdir10", Array.factory(DataType.FLOAT, dataShape, wdir10Data));
        System.out.println("风速风向计算完成");

        createWindNcFile(windVarMetadata, firstWindNc, croppedDataMap, croppedLatsArray, croppedLonsArray, timeSize, latSize, lonSize, outputPath);
    }

    // ==================== SST处理 ====================
    /**
     * 处理SST文件
     * @param fileMap  输入文件配置（路径->变量名）
     * @param outputPath 输出文件路径
     */
    public static void processSst(Map<String, String> fileMap, String outputPath) throws Exception {
        processFileGroup("SST", fileMap, () -> {
            try {
                innerProcessSst(fileMap, outputPath);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }

    private static void innerProcessSst(Map<String, String> fileMap, String outputPath) throws Exception {
        Map<String, Array> croppedDataMap = new HashMap<>();
        double[] lats = null;
        double[] lons = null;
        int latSize = -1;
        int lonSize = -1;
        Array croppedLatsArray = null;
        Array croppedLonsArray = null;
        NetcdfFile refNc = null;
        String validVarName = null;

        for (Map.Entry<String, String> entry : fileMap.entrySet()) {
            String filePath = entry.getKey();
            String varName = entry.getValue();

            if (!new File(filePath).exists()) {
                System.out.println("文件不存在，跳过：" + filePath);
                continue;
            }

            System.out.println("处理SST文件：" + filePath + "，变量：" + varName);
            validVarName = varName;

            NetcdfFile ncFile = NetcdfFile.open(filePath);
            if (refNc == null) refNc = ncFile;

            Variable var = ncFile.findVariable(varName);
            Variable latVar = ncFile.findVariable("latitude");
            Variable lonVar = ncFile.findVariable("longitude");
            validateVariables(var, latVar, lonVar, varName);

            int[] shape = var.getShape();
            boolean is3DWithSingleTime = (shape.length == 3 && shape[0] == 1);
            if (!(shape.length == 2 || is3DWithSingleTime)) {
                throw new RuntimeException("SST变量应为2维（lat, lon）或3维（time=1, lat, lon），实际维度：" + shape.length);
            }

            if (lats == null || lons == null) {
                lats = readLatLonArray(latVar);
                lons = readLatLonArray(lonVar);
                int[] latIndices = getRangeIndices(lats, MIN_LAT, MAX_LAT);
                int[] lonIndices = getRangeIndices(lons, MIN_LON, MAX_LON);
                if (latIndices == null || lonIndices == null) {
                    throw new RuntimeException("裁剪范围超出数据范围：" + filePath);
                }
                latSize = latIndices[1];
                lonSize = lonIndices[1];
                System.out.printf("SST裁剪后维度：纬度%d个点，经度%d个点%n", latSize, lonSize);
                croppedLatsArray = cropVariable(latVar, new int[]{latIndices[0]}, new int[]{latSize});
                croppedLonsArray = cropVariable(lonVar, new int[]{lonIndices[0]}, new int[]{lonSize});
            }

            Array croppedArray;
            if (shape.length == 3) {
                int[] origin = new int[]{0, getRangeIndices(lats, MIN_LAT, MAX_LAT)[0], getRangeIndices(lons, MIN_LON, MAX_LON)[0]};
                int[] cropShape = new int[]{1, latSize, lonSize};
                Array tempArray = cropVariable(var, origin, cropShape);
                croppedArray = tempArray.reduce();
            } else {
                int[] origin = new int[]{getRangeIndices(lats, MIN_LAT, MAX_LAT)[0], getRangeIndices(lons, MIN_LON, MAX_LON)[0]};
                int[] cropShape = new int[]{latSize, lonSize};
                croppedArray = cropVariable(var, origin, cropShape);
            }
            croppedDataMap.put(varName, croppedArray);

            ncFile.close();
            System.out.println("SST裁剪完成");
        }

        if (refNc != null && validVarName != null) {
            createSstNcFile(refNc, croppedDataMap, croppedLatsArray, croppedLonsArray, latSize, lonSize, outputPath, validVarName);
        } else {
            System.out.println("SST无有效文件，不生成输出");
        }
    }

    // ==================== 通用工具方法 ====================
    /**
     * 通用文件组处理逻辑
     * @param groupName  文件组名称（日志用）
     * @param fileMap    输入文件配置
     * @param process    具体处理逻辑
     */
    private static void processFileGroup(String groupName, Map<String, String> fileMap, Runnable process) {
        boolean hasValidFile = false;
        for (String filePath : fileMap.keySet()) {
            if (new File(filePath).exists()) {
                hasValidFile = true;
                break;
            }
        }

        if (hasValidFile) {
            System.out.println("\n=== 开始处理" + groupName + "文件 ===");
            try {
                process.run();
            } catch (Exception e) {
                System.err.println(groupName + "文件处理失败：" + e.getMessage());
                e.printStackTrace();
            }
        } else {
            System.out.println("\n=== " + groupName + "文件组所有文件均不存在，跳过处理 ===");
        }
    }

    /**
     * 创建带时间维度的单变量NC文件
     */
    private static void createSingleVarNcFile(NetcdfFile refNc, Map<String, Array> dataMap,
                                              double[] origLats, double[] origLons,
                                              int timeSize, int latSize, int lonSize,
                                              String outputPath, String varName, boolean hasTime) throws IOException {
        NetcdfFileWriter writer = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf4, outputPath);
        Group rootGroup = null;

        try {
            List<Dimension> dataDims = new ArrayList<>();
            if (hasTime) {
                Dimension timeDim = writer.addDimension(rootGroup, "time", timeSize);
                dataDims.add(timeDim);
            }
            Dimension latDim = writer.addDimension(rootGroup, "latitude", latSize);
            Dimension lonDim = writer.addDimension(rootGroup, "longitude", lonSize);
            dataDims.add(latDim);
            dataDims.add(lonDim);

            Variable latVar = writer.addVariable(rootGroup, "latitude",
                    refNc.findVariable("latitude").getDataType(), "latitude");
            Variable lonVar = writer.addVariable(rootGroup, "longitude",
                    refNc.findVariable("longitude").getDataType(), "longitude");
            copyVariableAttributes(refNc.findVariable("latitude"), latVar);
            copyVariableAttributes(refNc.findVariable("longitude"), lonVar);

            Variable dataVar = writer.addVariable(rootGroup, varName,
                    refNc.findVariable(varName).getDataType(), dataDims);
            copyVariableAttributes(refNc.findVariable(varName), dataVar);

            for (Attribute attr : refNc.getGlobalAttributes()) {
                writer.addGlobalAttribute(attr);
            }

            writer.create();
            int[] latIndices = getRangeIndices(origLats, MIN_LAT, MAX_LAT);
            int[] lonIndices = getRangeIndices(origLons, MIN_LON, MAX_LON);
            writer.write(latVar, cropVariable(refNc.findVariable("latitude"), new int[]{latIndices[0]}, new int[]{latSize}));
            writer.write(lonVar, cropVariable(refNc.findVariable("longitude"), new int[]{lonIndices[0]}, new int[]{lonSize}));
            writer.write(dataVar, dataMap.get(varName));

            System.out.println(varName + "文件创建成功：" + outputPath);

        } catch (InvalidRangeException e) {
            throw new RuntimeException(e);
        } finally {
            writer.close();
            refNc.close();
        }
    }

    /**
     * 创建风场NC文件
     */
    private static void createWindNcFile(Map<String, Variable> windVarMetadata, NetcdfFile firstWindNc,
                                         Map<String, Array> dataMap, Array croppedLatsArray, Array croppedLonsArray,
                                         int timeSize, int latSize, int lonSize, String outputPath) throws IOException {
        NetcdfFileWriter writer = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf4, outputPath);
        Group rootGroup = null;

        try {
            Dimension timeDim = writer.addDimension(rootGroup, "time", timeSize);
            Dimension latDim = writer.addDimension(rootGroup, "latitude", latSize);
            Dimension lonDim = writer.addDimension(rootGroup, "longitude", lonSize);
            List<Dimension> dataDims = Arrays.asList(timeDim, latDim, lonDim);

            Variable refLatVar = firstWindNc.findVariable("latitude");
            Variable refLonVar = firstWindNc.findVariable("longitude");
            Variable latVar = writer.addVariable(rootGroup, "latitude", refLatVar.getDataType(), "latitude");
            Variable lonVar = writer.addVariable(rootGroup, "longitude", refLonVar.getDataType(), "longitude");
            copyVariableAttributes(refLatVar, latVar);
            copyVariableAttributes(refLonVar, lonVar);

            Map<String, Variable> varMap = new HashMap<>();
            for (String name : Arrays.asList("u10", "v10")) {
                Variable refVar = windVarMetadata.get(name);
                Variable newVar = writer.addVariable(rootGroup, name, refVar.getDataType(), dataDims);
                copyVariableAttributes(refVar, newVar);
                varMap.put(name, newVar);
            }

            Variable wspd10Var = writer.addVariable(rootGroup, "wspd10", DataType.FLOAT, dataDims);
            wspd10Var.addAttribute(new Attribute("long_name", "10米高度风速"));
            wspd10Var.addAttribute(new Attribute("units", "m/s"));
            Variable wdir10Var = writer.addVariable(rootGroup, "wdir10", DataType.FLOAT, dataDims);
            wdir10Var.addAttribute(new Attribute("long_name", "10米高度风向"));
            wdir10Var.addAttribute(new Attribute("units", "degree"));
            varMap.put("wspd10", wspd10Var);
            varMap.put("wdir10", wdir10Var);

            for (Attribute attr : firstWindNc.getGlobalAttributes()) {
                writer.addGlobalAttribute(attr);
            }

            writer.create();
            writer.write(latVar, croppedLatsArray);
            writer.write(lonVar, croppedLonsArray);
            for (Map.Entry<String, Array> entry : dataMap.entrySet()) {
                Variable destVar = varMap.get(entry.getKey());
                if (destVar != null) writer.write(destVar, entry.getValue());
            }

            System.out.println("风场文件创建成功：" + outputPath);

        } catch (InvalidRangeException e) {
            throw new RuntimeException(e);
        } finally {
            writer.close();
            firstWindNc.close();
        }
    }

    /**
     * 创建SST NC文件
     */
    private static void createSstNcFile(NetcdfFile refNc, Map<String, Array> dataMap,
                                        Array croppedLatsArray, Array croppedLonsArray,
                                        int latSize, int lonSize, String outputPath, String varName) throws IOException {
        NetcdfFileWriter writer = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf4, outputPath);
        Group rootGroup = null;

        try {
            Dimension latDim = writer.addDimension(rootGroup, "latitude", latSize);
            Dimension lonDim = writer.addDimension(rootGroup, "longitude", lonSize);
            List<Dimension> dataDims = Arrays.asList(latDim, lonDim);

            Variable latVar = writer.addVariable(rootGroup, "latitude",
                    refNc.findVariable("latitude").getDataType(), "latitude");
            Variable lonVar = writer.addVariable(rootGroup, "longitude",
                    refNc.findVariable("longitude").getDataType(), "longitude");
            copyVariableAttributes(refNc.findVariable("latitude"), latVar);
            copyVariableAttributes(refNc.findVariable("longitude"), lonVar);

            Variable sstVar = writer.addVariable(rootGroup, varName,
                    refNc.findVariable(varName).getDataType(), dataDims);
            copyVariableAttributes(refNc.findVariable(varName), sstVar);

            for (Attribute attr : refNc.getGlobalAttributes()) {
                writer.addGlobalAttribute(attr);
            }

            writer.create();
            writer.write(latVar, croppedLatsArray);
            writer.write(lonVar, croppedLonsArray);
            writer.write(sstVar, dataMap.get(varName));

            System.out.println("SST文件创建成功：" + outputPath);

        } catch (InvalidRangeException e) {
            throw new RuntimeException(e);
        } finally {
            writer.close();
            refNc.close();
        }
    }

    // 基础校验与工具方法
    private static void validateVariables(Variable var, Variable latVar, Variable lonVar, String varName) {
        if (var == null) throw new RuntimeException("未找到变量：" + varName);
        if (latVar == null || lonVar == null) throw new RuntimeException("缺少经纬度变量");
        if (latVar.getRank() != 1 || lonVar.getRank() != 1) throw new RuntimeException("经纬度必须是1维变量");
    }

    private static double[] readLatLonArray(Variable var) throws IOException {
        Array array = var.read();
        DataType type = var.getDataType();
        if (type == DataType.DOUBLE) return (double[]) array.copyTo1DJavaArray();
        else if (type == DataType.FLOAT) {
            float[] floatArr = (float[]) array.copyTo1DJavaArray();
            double[] doubleArr = new double[floatArr.length];
            for (int i = 0; i < floatArr.length; i++) doubleArr[i] = floatArr[i];
            return doubleArr;
        } else throw new RuntimeException("经纬度类型不支持：" + type);
    }

    private static int[] getRangeIndices(double[] array, double min, double max) {
        int startIdx = -1, endIdx = -1;
        for (int i = 0; i < array.length; i++) {
            double val = array[i];
            if (val >= min && val <= max) {
                if (startIdx == -1) startIdx = i;
                endIdx = i;
            }
        }
        return (startIdx == -1 || endIdx == -1) ? null : new int[]{startIdx, endIdx - startIdx + 1};
    }

    private static Array cropVariable(Variable var, int[] origin, int[] shape) throws IOException, InvalidRangeException {
        Range[] ranges = new Range[origin.length];
        for (int i = 0; i < origin.length; i++) {
            ranges[i] = new Range(origin[i], origin[i] + shape[i] - 1);
        }
        return var.read(Arrays.asList(ranges));
    }

    private static void copyVariableAttributes(Variable srcVar, Variable destVar) {
        for (Attribute attr : srcVar.getAttributes()) {
            destVar.addAttribute(attr);
        }
    }
}